import logging
from typing import List

from llama_index.core.tools.types import DefaultToolFnSchema
from llama_index.core.workflow.handler import WorkflowHandler
from sqlalchemy.orm import Session

from app.models.ai import Tools, AIAgent, AIModel, AIApiKey
from app.models.base import BaseEnumEnabledDisabled
from sks.tools import get_tools, none_fun
from app.services.ai_services import AIService
from sks.agent.intelligent_agent import IntelligentAgent, AgentMeta
from app.db.session import get_db
from app.llama_index import message_to_ChatMessage
from app.llama_index.llms.openai import createOpenAI
from app.models import Message
from app.services.chat_services import getMessagesByChatId
from llama_index.core.tools import FunctionTool, ToolMetadata

#milvus集合名称+知识库ID
PREFIX_MILVUS_COLLECTION_NAME = 'aip_'

logger=logging.getLogger(__name__)


class BaseIntelligentAgent:
    tools:list=[]
    agent_tools:list=[] #作为工具的智能助手
    agent:IntelligentAgent
    ai_agent: AIAgent
    ai_model: AIModel
    ai_api: AIApiKey

    def __init__(self,db: Session,ai_agent:AIAgent,ai_model:AIModel, ai_api:AIApiKey,chat_id:int):
        self.db=db
        self.ai_agent=ai_agent
        self.ai_model=ai_model
        self.ai_api=ai_api
        self.chat_id=chat_id
        self.initTools()

        llm = createOpenAI(model=ai_model.model_name,
                           api_base=ai_api.api_base,
                           api_key=ai_api.api_key, temperature=0.7)

        self.chat_history=[]
        self.agent = IntelligentAgent(llm=llm,
                                      tools=self.tools,
                                      agents=self.agent_tools,
                                      system_prompt=ai_agent.system_prompt,
                                      chat_history=self.chat_history,
                                      chat_id=chat_id,
                                      metadata=AgentMeta(description=ai_agent.name,
                                                         tool_desc=ai_agent.tool_desc,
                                                         id=ai_agent.id,
                                                         kb_id=ai_agent.kb_id)
                                      )

    def initTools(self):
        # self.tools=[
        #     # tool_now_time,
        #     FunctionTool.from_defaults(tool_now_time),
        # ]
        self.tools=[]
        self.agent_tools=[]
        if self.ai_agent.tools and self.ai_agent.tools.strip()!='':
            tool_list=self.db.query(Tools).filter(Tools.status==BaseEnumEnabledDisabled.ENABLED.value.value,Tools.id.in_(self.ai_agent.tools.split(','))).all()
            _tools = get_tools()
            for item in tool_list:
                for tool in _tools:
                    if tool.metadata.name==item.tool_name:
                        self.tools.append(tool)
        if self.ai_agent.agent_ids and self.ai_agent.agent_ids.strip()!='':
            agent_list=self.db.query(AIAgent).filter(AIAgent.status==BaseEnumEnabledDisabled.ENABLED.value.value,AIAgent.id.in_(self.ai_agent.agent_ids.split(','))).all()
            for item in agent_list:
                ai_agent, ai_model, ai_api =AIService.getAILLMInfoByAIAgentId(self.db,item.id)
                iab=BaseIntelligentAgent(self.db,ai_agent=item,ai_model=ai_model,ai_api=ai_api,chat_id=self.chat_id)
                self.agent_tools.append(iab.agent)
        print(self.tools)

    def add_history(self,messageDTO:Message):
        _cm=message_to_ChatMessage(messageDTO)
        self.chat_history.append(_cm)

    def reset_history(self,messages:List[Message]):
        self.chat_history.clear()
        for messageDTO in messages:
            self.add_history(messageDTO)


class LlamaIndexFunctionAgentSession(BaseIntelligentAgent):
    agent:IntelligentAgent
    tools:list=[]
    def __init__(self,db: Session,chat_id:int):
        self.chat_id=chat_id
        self.db=db

        chat, ai_agent,ai_model, ai_api = AIService.getAILLMInfo(self.db,chat_id)
        super().__init__(db=db,ai_agent=ai_agent,ai_model=ai_model,ai_api=ai_api,chat_id=chat_id)
        # iab=BaseIntelligentAgent(db,ai_agent=ai_agent,ai_model=ai_model,ai_api=ai_api,chat_id=chat_id)
        # self.agent=iab.agent



    def run(self,messageDTO:Message):
        user_msg=messageDTO.content
        # logger.debug(f'{self.ctx}')
        # logger.debug(f'{self.ctx.to_dict()}')
        # response:WorkflowHandler=self.agent.chat(user_msg=user_msg,chat_id=self.chat_id)
        response:WorkflowHandler=self.agent.chat_with_tools(user_msg=user_msg,chat_id=self.chat_id)
        res=response
        logger.debug(f'{res}')
        # logger.debug(f'{response["response"]}')
        # logger.debug(f'{self.memory.get()}')
        return str(res)
class LlamaIndexFunctionAgentSessions:
    _sessions: map = {}

    def __init__(self):
        pass

    def get_session(self,chat_id) -> LlamaIndexFunctionAgentSession:
        #不换存,每次都用新的session
        # if chat_id in self._sessions:
        #     return self._sessions[chat_id]
        db: Session = next(get_db())
        session= LlamaIndexFunctionAgentSession(db=db,chat_id=chat_id)
        messages=getMessagesByChatId(db,chat_id)
        session.reset_history(messages=messages)
        #不换存,每次都用新的session
        # self._sessions[chat_id]=session
        return session


lifaSessions = LlamaIndexFunctionAgentSessions()
